
# Copyright 2009 David Lee Miller
#
# This file is part of POSSM (Python Open Source Scanning Microscopy).
#
# POSSM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POSSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License         
# along with POSSM.  If not, see <http://www.gnu.org/licenses/>.
#                                                                           
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from client import sys, os, gc
from client import matplotlib,rcParams,cm,Toolbar,FigureCanvasWxAgg,NavigationToolbar2WxAgg
from client import Figure,nx,RectangleSelector,_load_bitmap

from client import time,wx
from client import n #the numpy package
#from plotarray import *
#from wxmatplot import PlotPanel
TIMER_ID = wx.NewId()

class MyNavigationToolbar(NavigationToolbar2WxAgg):
	"""
	Extend the default wx toolbar with your own event handlers
	"""
	ON_IV = wx.NewId()
	ON_GRAY = wx.NewId()
	def __init__(self, canvas, cankill):
		NavigationToolbar2WxAgg.__init__(self, canvas)
		self.a = None
		self.fig = None
		self.transiv = None
		self.ivsw = 0
		self.graysw = 0
		self.setcolor = cm.jet
		# for simplicity I'm going to reuse a bitmap from wx, you'll
		# probably want to add your own.
		#self.AddSimpleTool(self.ON_IV, _load_bitmap('stock_refresh.xpm'),
					#'IV', 'Toggle IV Display')
		#EVT_TOOL(self, self.ON_IV, self._on_custom)
		
		self.AddSimpleTool(self.ON_GRAY, _load_bitmap('stock_refresh.xpm'),
					'Click me', 'Toggle Grayscale')
		EVT_TOOL(self, self.ON_GRAY, self._on_gray)
	
	def get_tools(self,im,fig,canvas):
		self.im = im
		self.fig = fig
		#self.package = package
		self.canvas = canvas
	
	#def _on_clear(self, evt):
		#if self.fig and self.a:
			#self.a.clear()
			#self.a.imshow(toPylab(self.package.IMAGEDATA), cmap=cm.jet)
			#self.canvas.draw()
			#EVT_TOOL(self, self.ON_CUSTOM, self._on_custom)
	
	def _on_gray(self, evt):
		if self.graysw == 0:
			self.graysw = 1
			#print dir(self.a)
			#print dir(self.im)
			#print dir(self.fig)
			#print dir(cm)
			self.setcolor = cm.gray
		elif self.graysw == 1:
			self.graysw = 2
			self.setcolor = cm.gist_stern
		elif self.graysw == 2:
			self.graysw = 3
			self.setcolor = cm.hot
		elif self.graysw == 3:
			self.graysw = 4
			self.setcolor = cm.gist_earth
		elif self.graysw == 4:
			self.graysw = 5
			self.setcolor = cm.jet
		elif self.graysw == 5:
			self.graysw = 0
			self.setcolor = cm.jet_r
		
		self.im.set_cmap(self.setcolor)
		#self.im.set_colorbar(self.im)
		self.canvas.draw()
		evt.Skip()

class ArrangeImages(object):
	MASTER = n.zeros((1000,1000))
	
	def _AddImage(self,package):
		xoff = package.XOFF[0]
		yoff = package.YOFF[0]
		print '1'
		dac_step = package.DAC_STEP[0]
		dac_avgs = package.DAC_AVGS[0]
		scan_y = package.SCAN_Y[0]
		scan_x = package.SCAN_X[0]
		print '2'
		self.TOTAL_OFF_VOLTS = 20
		self.TOTAL_IMG_VOLTS = 20
		total_range = self.TOTAL_OFF_VOLTS + self.TOTAL_IMG_VOLTS
		possible_ndac = int(total_range * (65536./self.TOTAL_IMG_VOLTS))
		actual_npts = self.MASTER.shape[0]
		print '3',actual_npts,scan_x,dac_step,dac_avgs
		npt_per_dac = (float(actual_npts)/possible_ndac)
		dac_x = scan_x * dac_step * dac_avgs
		dac_y = scan_y * dac_step * dac_avgs
		print '4'
		scaled_x = npt_per_dac * dac_x
		scaled_y = npt_per_dac * dac_y
		x_scale = int(round(scan_x / scaled_x))
		y_scale = int(round(y / scaled_y))
		print '5',x_scale,y_scale,package.IMAGEDATA[0][0]
		new_img = toPylab(package.IMAGEDATA[0][0][::y_scale][::x_scale])
		startx = int(round((float(xoff)/total_range)*actual_npts)+actual_npts/2.)
		print '6'
		starty = int(round((float(yoff)/total_range)*actual_npts)+actual_npts/2.)
		endx = startx + new_img.shape[1]
		endy = starty + new_img.shape[0]
		print '7'
		self.MASTER[starty:endy,startx:endx] = new_img

class PiezoMap(wx.Frame):
	MASTER = n.zeros((1000,1000))
	def __init__(self,parent,name,data):
		wx.Frame.__init__(self, parent, -1, name)
		
		self.fig = Figure((7,6), 60)
		#self.fig = Figure(None, None)
		self.canvas = FigureCanvasWxAgg(self, -1, self.fig)
		self.toolbar = MyNavigationToolbar(self.canvas, True)
		self.toolbar.Realize()
		self.x = None
		self.y = None
		self.ivsub = None
		self.coords = None
		self.timelast = 0
		self.update = 0
		#self.package = package
		self.name = name
		
		# On Windows, default frame size behaviour is incorrect
		# you don't need this under Linux
		tw, th = self.toolbar.GetSizeTuple()
		fw, fh = self.canvas.GetSizeTuple()
		self.toolbar.SetSize(wx.Size(fw, th))
		
		#self.init_plot_data()
		# Initialise the timer - wxPython requires this to be connected to
		# the receiving event handler
		self.t = wx.Timer(self, TIMER_ID)
		self.t.Start(10000)
		self.txt = None
		# Create a figure manager to manage things
		
		# Now put all into a sizer
		sizer = wx.BoxSizer(wx.VERTICAL)
		# This way of adding to sizer allows resizing
		sizer.Add(self.canvas, 1, wx.LEFT|wx.TOP|wx.GROW)
		# Best to allow the toolbar to resize!
		sizer.Add(self.toolbar, 0, wx.GROW)
		self.SetSizer(sizer)
		self.Fit()
		#self.canvas.Bind(wx.EVT_RIGHT_DOWN, self.OnRightDown)
		self.Bind(wx.EVT_TIMER, self.onTimer, id=TIMER_ID)
		self.Bind(wx.EVT_CLOSE, self.onClose)
		
		# Note that event is a MplEvent
		self.canvas.mpl_connect('motion_notify_event', self.UpdateStatusBar)
		self.canvas.Bind(wx.EVT_ENTER_WINDOW, self.ChangeCursor)
		
		self.statusBar = wx.StatusBar(self, -1)
		self.statusBar.SetFieldsCount(1)
		self.SetStatusBar(self.statusBar)
		
		self._Plot(data)
		self.Show()

	def _Plot(self,data):
		# jdh you can add a subplot directly from the fig rather than
		# the fig manager
		a = self.fig.add_subplot(111)
		self.im = a.imshow( data, cmap=cm.jet)
		#self.im = a.imshow( toPylab(self.Z), cmap=cm.jet)
		self.fig.colorbar(self.im)
		self.toolbar.get_tools(self.im,self.fig,self.canvas)
		#print dir(self.im)
	
	#def _Show(self,parent,name):
		#self.__init__(parent,name)
		#self._Plot()
		#self.Show()
	
	def ChangeCursor(self, event):
		self.canvas.SetCursor(wx.StockCursor(wx.CURSOR_ARROW))
	
	#def _AddImage(self,package):
		#xoff = package.XOFF[0]
		#yoff = package.YOFF[0]
		#print '1'
		#dac_step = package.DAC_STEP[0]
		#dac_avgs = package.DAC_AVGS[0]
		#height = package.HEIGHT[0]
		#width = package.WIDTH[0]
		#print '2'
		#self.TOTAL_OFF_VOLTS = 20
		#self.TOTAL_IMG_VOLTS = 20
		#total_range = self.TOTAL_OFF_VOLTS + self.TOTAL_IMG_VOLTS
		#possible_ndac = int(total_range * (65536./self.TOTAL_IMG_VOLTS))
		#actual_npts = self.MASTER.shape[0]
		#print '3',actual_npts,width,dac_step,dac_avgs
		#npt_per_dac = (float(actual_npts)/possible_ndac)
		#dac_width = width * dac_step * dac_avgs
		#dac_height = height * dac_step * dac_avgs
		#print '4'
		#scaled_width = npt_per_dac * dac_width
		#scaled_height = npt_per_dac * dac_height
		#width_scale = int(round(width / scaled_width))
		#height_scale = int(round(height / scaled_height))
		#print '5',width_scale,height_scale,package.IMAGEDATA
		#new_img = toPylab(package.IMAGEDATA[::height_scale,::width_scale])
		#startx = int(round((float(xoff)/total_range)*actual_npts)+actual_npts/2.)
		#print '6'
		#starty = int(round((float(yoff)/total_range)*actual_npts)+actual_npts/2.)
		#endx = startx + new_img.shape[1]
		#endy = starty + new_img.shape[0]
		#print '7'
		#self.MASTER[starty:endy,startx:endx] = new_img
		
		
	
	def UpdateStatusBar(self, event):
		if event.inaxes:
			self.x, self.y = event.xdata, event.ydata
			#if type(self.x) != float or type(self.y) != float:
			self.x = round(self.x)
			self.y = round(self.y)
			self.statusBar.SetStatusText(( "x= " + str(int(round(self.x))) +
							"  y=" +str(int(round(self.y))) ),
							0)
		#if self.ivsub:
			#self.coords = str((int(self.y),int(self.x)))
			#if self.coords in self.IV and time.time()-self.timelast > 0.075:
				##self.ivsub._drawOnChange(self.coords,self.IV[self.coords])
				#self.timelast = time.time()

	def onTimer(self, evt):
		#self.im.set_array(toPylab(self.Z))
		self.canvas.draw()

	def onClose(self, evt):
		self.t.Stop()
		evt.Skip()

	#def OnRightDown(self, evt):
		#"""right mouse button is pressed"""
		##print (int(round(self.x),int(round(self.y))))
		#if not self.ivsub:
			#self.ivsub = IVFigure(None,self.name + ' IV')
		#else:
			#coords = str((self.x,self.y))
			#self.ivsub._drawOnChange(coords,self.IV[self.IV.keys()[0]])